package com.duyus.ai.rag;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
public class PgVectorVectorStoreConfig {

//    @Resource
//    private MdDocumentLoader loveAppDocumentLoader;

    @Autowired
    @Qualifier("postgresDataSource") // 明确指定你要用哪个数据源
    private DataSource postgresDataSource;

    @Bean
    public VectorStore pgVectorVectorStore(@Qualifier("dashscopeEmbeddingModel") EmbeddingModel dashscopeEmbeddingModel) {
        // 指定数据库 可以原生操作数据库
        JdbcTemplate jdbcTemplate = new JdbcTemplate(postgresDataSource);
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // 可选，默认根据 embedding model 自动推断或 1536
                .distanceType(COSINE_DISTANCE)       // 可选，默认 COSINE_DISTANCE
                .indexType(HNSW)                     // 可选，默认 HNSW
                .initializeSchema(true)              // 可选，默认 false
                .schemaName("public")                // 可选，默认 public
                .vectorTableName("vector_store")     // 可选，默认 vector_store
                .maxDocumentBatchSize(10000)         // 可选，默认 10000
                .build();

//        // 加载文档并插入向量库
//        List<Document> documents = loveAppDocumentLoader.loadMarkdowns();
//        vectorStore.add(documents);

        return vectorStore;
    }
}